from torch import nn  
from common import *
import torch

class UnetGenerator(nn.Module):
    def __init__(self, channel=32, num_blocks=4,n_class=19):
        super().__init__()

        self.conv = nn.Conv2d(3, channel, [7, 7], padding=3)  # same 256,256
        self.conv1 = nn.Conv2d(channel, channel, [3, 3], stride=2, padding=1)  # same 128,128
        self.conv2 = nn.Conv2d(channel, channel * 2, [3, 3], padding=1)  # 128,128
        self.conv3 = nn.Conv2d(channel * 2, channel * 2, [3, 3], stride=2, padding=1)  # 64,64
        self.conv4 = nn.Conv2d(channel * 2, channel * 4, [3, 3], padding=1)  # 64,64

        self.resblock = nn.Sequential(*[ResBlock(channel * 4, channel * 4) for i in range(num_blocks)])

        self.conv5 = nn.Conv2d(channel * 4, channel * 2, [3, 3], padding=1)  # 64,64
        self.conv6 = nn.Conv2d(channel * 2, channel * 2, [3, 3], padding=1)  # 64,64
        self.conv7 = nn.Conv2d(channel * 2, channel, [3, 3], padding=1)  # 64,64
        self.conv8 = nn.Conv2d(channel, channel, [3, 3], padding=1)  # 64,64
        self.conv9 = nn.Conv2d(channel, 3, [7, 7], padding=3)  # 64,64

        self.leak_relu = nn.LeakyReLU(inplace=True)
        self.upsample = nn.UpsamplingBilinear2d(scale_factor=2)
        # self.upsample1 = nn.ConvTranspose2d(64,64,)
        # self.upsample2 = nn.ConvTranspose2d()
        # nn.UpsamplingNearest2d
        self.act = nn.Tanh()
        
        
        
        self.seresblock = nn.Sequential(*[ResBlock(channel * 4, channel * 4) for i in range(num_blocks)])

        self.seconv5 = nn.Conv2d(channel * 4, channel * 2, [3, 3], padding=1)  # 64,64
        self.sebn5 = nn.BatchNorm2d(channel * 2)
        self.seconv6 = nn.Conv2d(channel * 2, channel * 2, [3, 3], padding=1)  # 64,64
        self.sebn6 = nn.BatchNorm2d(channel * 2)
        self.seconv7 = nn.Conv2d(channel * 2, channel, [3, 3], padding=1)  # 64,64
        self.sebn7 = nn.BatchNorm2d(channel)
        self.seconv8 = nn.Conv2d(channel, channel, [3, 3], padding=1)  # 64,64
        self.sebn8 = nn.BatchNorm2d(channel)
        self.seconv9 = nn.Conv2d(channel, n_class, [7, 7], padding=3)  # 64,64
        self.relu = nn.ReLU(inplace=True)
        
        
        
        
          # facelandmark
        self.fconv4_1 = InvertedResidual(128,128,2,False,2)

        # 第5层卷积层，Bottleneck结构
        # 重复6次
        self.fconv5_1 = InvertedResidual(128,128,1,False,4)
        self.fblock5_2 = InvertedResidual(128,128,1,True,4)
        self.fblock5_3 = InvertedResidual(128,128,1,True,4)
        self.fblock5_4 = InvertedResidual(128,128,1,True,4)
        self.fblock5_5 = InvertedResidual(128,128,1,True,4)
        self.fblock5_6 = InvertedResidual(128,128,1,True,4)

        # 第6层卷积层，Bottleneck结构
        # 重复1次
        self.fconv6_1 = InvertedResidual(128,16,1,False,2)

        # 第7层卷积层，重复1次
        self.fconv7 = conv_bn(16,32,3,2)
        # 第8层卷积层，重复1次
        self.fconv8 = nn.Conv2d(32,64,7,1,0)
        self.fbn8 = nn.BatchNorm2d(64)


        self.favg_pool1 = nn.AvgPool2d(14)
        self.favg_pool2 = nn.AvgPool2d(7)
        self.favg_pool3 = nn.AvgPool2d(5)

        self.ffc = nn.Linear(448,196)

        self.frelu = nn.ReLU(inplace=True)
        
        

    def forward(self, mem):
        landmark = None
        gan = None
        seg = None
        
        if 'gan' in mem.keys():
            inputs_gan = mem['gan']['data']
            gan = self.forward_gan(inputs_gan)
        if 'seg' in mem.keys():
            inputs_seg = mem['seg']['data']
            seg = self.forward_seg(inputs_seg)
        if 'facelandmark' in mem.keys():
            inputs_facelandmark = mem['facelandmark']['data']
            landmark = self.forward_facelandmark(inputs_facelandmark)
        
        return gan,landmark,seg
        
    def forward_gan(self,inputs):
        x0 = self.conv(inputs)
        x0 = self.leak_relu(x0)  # 256, 256, 32

        x1 = self.conv1(x0)
        x1 = self.leak_relu(x1)
        x1 = self.conv2(x1)
        x1 = self.leak_relu(x1)  # 128, 128, 64

        x2 = self.conv3(x1)
        x2 = self.leak_relu(x2)
        x2 = self.conv4(x2)
        x2 = self.leak_relu(x2)  # 64, 64, 128

        x2 = self.resblock(x2)  # 64, 64, 128

        x2 = self.conv5(x2)
        x2 = self.leak_relu(x2)  # 64, 64, 64
        x3 = self.upsample(x2)
        x3 = self.conv6(x3 )
        x3 = self.leak_relu(x3)
        
        x3 = self.conv7(x3)
        x3 = self.leak_relu(x3)  # 128, 128, 32
        x4 = self.upsample(x3)
        x4 = self.conv8(x4 )
        x4 = self.leak_relu(x4)
        
        x5 = self.conv9(x4)  # 256, 256, 32
        x5 = self.leak_relu(x5)
        
        result = self.act(x5)

        return result
    
    def forward_seg(self,x):
        start_b = x.shape[0]
        
        x0 = self.conv(x)
        # x0 = self.bn(x0)
        x0 = self.leak_relu(x0)  # 256, 256, 32

        x1 = self.conv1(x0)
        x1 = self.leak_relu(x1)
        # x1 = self.bn1(x1)
        x1 = self.conv2(x1)
        x1 = self.leak_relu(x1)  # 128, 128, 64'
        # x1 = self.bn2(x1)

        x2 = self.conv3(x1)
        x2 = self.leak_relu(x2)
        # x2 = self.bn3(x2)
        x2 = self.conv4(x2)
        x2 = self.leak_relu(x2)  # 64, 64, 128
        # x2 = self.bn4(x2)
        
        x2 = self.seresblock(x2)  # 64, 64, 128
        x2 = self.seconv5(x2)
        x2 = self.relu(x2)  # 64, 64, 64
        x2 = self.sebn5(x2)

        x3 = self.upsample(x2)
        x3 = self.seconv6(x3 + x1)
        x3 = self.relu(x3)
        x3 = self.sebn6(x3)
        x3 = self.seconv7(x3)
        x3 = self.relu(x3)  # 128, 128, 32
        x3 = self.sebn7(x3)

        x4 = self.upsample(x3)
        x4 = self.seconv8(x4 + x0)
        x4 = self.relu(x4)
        x4 = self.sebn8(x4)
        x4 = self.seconv9(x4)  # 256, 256, 32
        finish_b = x4.shape[0]
        # print(start_b,finish_b)
        if finish_b != start_b:
            
            import pdb
            pdb.set_trace()
            
            
        return x4

    def forward_facelandmark(self,inputs):
            x0 = self.conv(inputs)
            x0 = self.leak_relu(x0)  # 256, 256, 32

            x1 = self.conv1(x0)
            x1 = self.leak_relu(x1)
            x1 = self.conv2(x1)
            x1 = self.leak_relu(x1)  # 128, 128, 64'

            x2 = self.conv3(x1)
            x2 = self.leak_relu(x2)
            x2 = self.conv4(x2)
            x2 = self.leak_relu(x2)  # 64, 64, 128
            
            x2= self.resblock(x2)
            
            x41=self.fconv4_1(x2)
            x51=self.fconv5_1(x41)
            x52=self.fblock5_2(x51)
            x53=self.fblock5_3(x52)
            x54=self.fblock5_4(x53)
            x55=self.fblock5_5(x54)
            x56=self.fblock5_6(x55)
            
            x61=self.fconv6_1(x56)
            
            output1= self.favg_pool1(x61)
            output1 = output1.view(output1.size(0),-1)
            
            x7=self.fconv7(x61)
            output2 = self.favg_pool2(x7)
            output2 = output2.view(output2.size(0),-1)
            output3 = self.favg_pool3(self.frelu(self.fconv8(x7)))
            output3 = output3.view(output3.size(0),-1)

            multi_scale = torch.cat([output1,output2,output3],1)
            labdmarks = self.ffc(multi_scale)
            
            return labdmarks





class Discriminator(nn.Module):
    def __init__(self, in_channels=3, out_channels=1, features=[32, 64, 128]):
        super().__init__()
        
        self.model = nn.Sequential(
            #k3n32s2
            Block(in_channels, features[0], kernel_size=3, stride=2, padding=1),
            #k3n32s1
            Block(features[0], features[0], kernel_size=3, stride=1, padding=1),

            #k3n64s2
            Block(features[0], features[1], kernel_size=3, stride=2, padding=1),
            #k3n64s1
            Block(features[1], features[1], kernel_size=3, stride=1, padding=1),

            #k3n128s2
            Block(features[1], features[2], kernel_size=3, stride=2, padding=1),
            #k3n128s1
            Block(features[2], features[2], kernel_size=3, stride=1, padding=1),

            #k1n1s1
            Block(features[2], out_channels, kernel_size=1, stride=1, padding=0,activate=False)
        )

    def forward(self, x):
        x = self.model(x)
        return x
    
    
class BiSeNet(nn.Module):
    def __init__(self, n_classes, *args, **kwargs):
        super(BiSeNet, self).__init__()
        self.cp = ContextPath()
        ## here self.sp is deleted
        self.ffm = FeatureFusionModule(256, 256)
        self.conv_out = BiSeNetOutput(256, 256, n_classes)
        self.conv_out16 = BiSeNetOutput(128, 64, n_classes)
        self.conv_out32 = BiSeNetOutput(128, 64, n_classes)
        self.init_weight()

    def forward(self, x):
        H, W = x.size()[2:]
        feat_res8, feat_cp8, feat_cp16 = self.cp(x)  # here return res3b1 feature
        feat_sp = feat_res8  # use res3b1 feature to replace spatial path feature
        feat_fuse = self.ffm(feat_sp, feat_cp8)

        feat_out = self.conv_out(feat_fuse)
        feat_out16 = self.conv_out16(feat_cp8)
        feat_out32 = self.conv_out32(feat_cp16)

        feat_out = F.interpolate(feat_out, (H, W), mode='bilinear', align_corners=True)
        feat_out16 = F.interpolate(feat_out16, (H, W), mode='bilinear', align_corners=True)
        feat_out32 = F.interpolate(feat_out32, (H, W), mode='bilinear', align_corners=True)
        return feat_out, feat_out16, feat_out32

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)

    def get_params(self):
        wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], []
        for name, child in self.named_children():
            child_wd_params, child_nowd_params = child.get_params()
            if isinstance(child, FeatureFusionModule) or isinstance(child, BiSeNetOutput):
                lr_mul_wd_params += child_wd_params
                lr_mul_nowd_params += child_nowd_params
            else:
                wd_params += child_wd_params
                nowd_params += child_nowd_params
        return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params



class UnetGenerator5(nn.Module):
    def __init__(self, input_channel = 3,channel=16, num_blocks=2,n_class=19):
        super(UnetGenerator5,self).__init__()

        self.conv = DepthWiseConv(input_channel, channel, [7, 7], padding=3)  # same 256,256
        self.conv1 = DepthWiseConv(channel, channel, [3, 3], stride=2, padding=1)  # same 128,128
        self.conv2 = DepthWiseConv(channel, channel * 2, [3, 3], padding=1)  # 128,128
        self.conv3 = DepthWiseConv(channel * 2, channel * 2, [3, 3], stride=2, padding=1)  # 64,64
        self.conv4 = DepthWiseConv(channel * 2, channel * 4, [3, 3], padding=1)  # 64,64

        self.resblock = nn.Sequential(*[ResBlockDepwise(channel * 4, channel * 4) for i in range(num_blocks)])

        self.conv5 = DepthWiseConv(channel * 4, channel * 2, [3, 3], padding=1)  # 64,64
        self.conv6 = DepthWiseConv(channel * 2, channel * 2, [3, 3], padding=1)  # 64,64
        self.conv7 = DepthWiseConv(channel * 2, channel, [3, 3], padding=1)  # 64,64
        self.conv8 = DepthWiseConv(channel, channel, [3, 3], padding=1)  # 64,64
        self.conv9 = DepthWiseConv(channel, 3, [7, 7], padding=3)  # 64,64
        self.upsample = nn.UpsamplingBilinear2d(scale_factor=2)
        self.act = nn.Tanh()
        
        
        
        
        self.seresblock = nn.Sequential(*[ResBlock(channel * 4, channel * 4) for i in range(num_blocks)])

        self.seconv5 = nn.Conv2d(channel * 4, channel * 2, [3, 3], padding=1)  # 64,64
        self.sebn5 = nn.BatchNorm2d(channel * 2)
        self.seconv6 = nn.Conv2d(channel * 2, channel * 2, [3, 3], padding=1)  # 64,64
        self.sebn6 = nn.BatchNorm2d(channel * 2)
        self.seconv7 = nn.Conv2d(channel * 2, channel, [3, 3], padding=1)  # 64,64
        self.sebn7 = nn.BatchNorm2d(channel)
        self.seconv8 = nn.Conv2d(channel, channel, [3, 3], padding=1)  # 64,64
        self.sebn8 = nn.BatchNorm2d(channel)
        self.seconv9 = nn.Conv2d(channel, n_class, [7, 7], padding=3)  # 64,64
        self.relu = nn.ReLU(inplace=True)
        
        
        
        
          # facelandmark
        self.fconv4_1 = InvertedResidual(128,128,2,False,2)

        # 第5层卷积层，Bottleneck结构
        # 重复6次
        self.fconv5_1 = InvertedResidual(128,128,1,False,4)
        self.fblock5_2 = InvertedResidual(128,128,1,True,4)
        self.fblock5_3 = InvertedResidual(128,128,1,True,4)
        self.fblock5_4 = InvertedResidual(128,128,1,True,4)
        self.fblock5_5 = InvertedResidual(128,128,1,True,4)
        self.fblock5_6 = InvertedResidual(128,128,1,True,4)

        # 第6层卷积层，Bottleneck结构
        # 重复1次
        self.fconv6_1 = InvertedResidual(128,16,1,False,2)

        # 第7层卷积层，重复1次
        self.fconv7 = conv_bn(16,32,3,2)
        # 第8层卷积层，重复1次
        self.fconv8 = nn.Conv2d(32,64,7,1,0)
        self.fbn8 = nn.BatchNorm2d(64)


        self.favg_pool1 = nn.AvgPool2d(14)
        self.favg_pool2 = nn.AvgPool2d(7)
        self.favg_pool3 = nn.AvgPool2d(5)

        self.ffc = nn.Linear(448,196)

        self.frelu = nn.ReLU(inplace=True)
        
        self.leak_relu = nn.LeakyReLU(inplace=True)

    def forward_gan(self, inputs):
        x0 = self.conv(inputs)
        x1 = self.conv1(x0)
        x1 = self.conv2(x1)
        x2 = self.conv3(x1)
        x2 = self.conv4(x2)
        x2 = self.resblock(x2)  # 64, 64, 128
        x2 = self.conv5(x2)
        x3 = self.upsample(x2)
        x3 = self.conv6(x3)
        x3 = self.conv7(x3)
        x4 = self.upsample(x3)
        x4 = self.conv8(x4)
        x4 = self.conv9(x4)  # 256, 256, 32
        result = self.act(x4)

        return torch.clamp(result,-0.999999,0.999999)
    def forward_seg(self,x):
        start_b = x.shape[0]
        
        x0 = self.conv(x)
        # x0 = self.bn(x0)
        x0 = self.leak_relu(x0)  # 256, 256, 32

        x1 = self.conv1(x0)
        x1 = self.leak_relu(x1)
        # x1 = self.bn1(x1)
        x1 = self.conv2(x1)
        x1 = self.leak_relu(x1)  # 128, 128, 64'
        # x1 = self.bn2(x1)

        x2 = self.conv3(x1)
        x2 = self.leak_relu(x2)
        # x2 = self.bn3(x2)
        x2 = self.conv4(x2)
        x2 = self.leak_relu(x2)  # 64, 64, 128
        # x2 = self.bn4(x2)
        
        x2 = self.seresblock(x2)  # 64, 64, 128
        x2 = self.seconv5(x2)
        x2 = self.relu(x2)  # 64, 64, 64
        x2 = self.sebn5(x2)

        x3 = self.upsample(x2)
        x3 = self.seconv6(x3 + x1)
        x3 = self.relu(x3)
        x3 = self.sebn6(x3)
        x3 = self.seconv7(x3)
        x3 = self.relu(x3)  # 128, 128, 32
        x3 = self.sebn7(x3)

        x4 = self.upsample(x3)
        x4 = self.seconv8(x4 + x0)
        x4 = self.relu(x4)
        x4 = self.sebn8(x4)
        x4 = self.seconv9(x4)  # 256, 256, 32
        finish_b = x4.shape[0]
        # print(start_b,finish_b)
        if finish_b != start_b:
            
            import pdb
            pdb.set_trace()
            
            
        return x4

    def forward_facelandmark(self,inputs):
            x0 = self.conv(inputs)
            x0 = self.leak_relu(x0)  # 256, 256, 32

            x1 = self.conv1(x0)
            x1 = self.leak_relu(x1)
            x1 = self.conv2(x1)
            x1 = self.leak_relu(x1)  # 128, 128, 64'

            x2 = self.conv3(x1)
            x2 = self.leak_relu(x2)
            x2 = self.conv4(x2)
            x2 = self.leak_relu(x2)  # 64, 64, 128
            
            x2= self.resblock(x2)
            
            x41=self.fconv4_1(x2)
            x51=self.fconv5_1(x41)
            x52=self.fblock5_2(x51)
            x53=self.fblock5_3(x52)
            x54=self.fblock5_4(x53)
            x55=self.fblock5_5(x54)
            x56=self.fblock5_6(x55)
            
            x61=self.fconv6_1(x56)
            
            output1= self.favg_pool1(x61)
            output1 = output1.view(output1.size(0),-1)
            
            x7=self.fconv7(x61)
            output2 = self.favg_pool2(x7)
            output2 = output2.view(output2.size(0),-1)
            output3 = self.favg_pool3(self.frelu(self.fconv8(x7)))
            output3 = output3.view(output3.size(0),-1)

            multi_scale = torch.cat([output1,output2,output3],1)
            labdmarks = self.ffc(multi_scale)
            
            return labdmarks

    def forward(self, mem):
        landmark = None
        gan = None
        seg = None
        
        if 'gan' in mem.keys():
            inputs_gan = mem['gan']['data']
            gan = self.forward_gan(inputs_gan)
        if 'seg' in mem.keys():
            inputs_seg = mem['seg']['data']
            seg = self.forward_seg(inputs_seg)
        if 'facelandmark' in mem.keys():
            inputs_facelandmark = mem['facelandmark']['data']
            landmark = self.forward_facelandmark(inputs_facelandmark)
        
        return gan,landmark,seg
    
    
    
class UnetGenerator10(nn.Module):
    def __init__(self, input_channel = 3,channel=16, num_blocks=2,n_class=19):
        super().__init__()

        self.conv = nn.Conv2d(input_channel, channel, [7, 7], padding=3)  # same 256,256
        self.bn = nn.BatchNorm2d(channel)
        self.conv1 = nn.Conv2d(channel, channel, [3, 3], stride=2, padding=1)  # same 128,128
        self.bn1 = nn.BatchNorm2d(channel)
        self.conv2 = nn.Conv2d(channel, channel * 2, [3, 3], padding=1)  # 128,128
        self.bn2 = nn.BatchNorm2d(channel * 2)
        self.conv3 = nn.Conv2d(channel * 2, channel * 2, [3, 3], stride=2, padding=1)  # 64,64
        self.bn3 = nn.BatchNorm2d(channel * 2)
        self.conv4 = nn.Conv2d(channel * 2, channel * 4, [3, 3], padding=1)  # 64,64
        self.bn4 = nn.BatchNorm2d(channel * 4)
        # gan
        self.resblock = nn.Sequential(*[ResBlock(channel * 4, channel * 4) for i in range(num_blocks)])

        self.conv5 = nn.Conv2d(channel * 4, channel * 2, [3, 3], padding=1)  # 64,64
        self.bn5 = nn.BatchNorm2d(channel * 2)
        self.conv6 = nn.Conv2d(channel * 2, channel * 2, [3, 3], padding=1)  # 64,64
        self.bn6 = nn.BatchNorm2d(channel * 2)
        self.conv7 = nn.Conv2d(channel * 2, channel, [3, 3], padding=1)  # 64,64
        self.bn7 = nn.BatchNorm2d(channel)
        self.conv8 = nn.Conv2d(channel, channel, [3, 3], padding=1)  # 64,64
        self.bn8 = nn.BatchNorm2d(channel)
        self.conv9 = nn.Conv2d(channel, 3, [7, 7], padding=3)  # 64,64
        self.bn9 = nn.BatchNorm2d(3)

        self.leak_relu = nn.LeakyReLU(0.0625,inplace=True)
        self.upsample = nn.UpsamplingNearest2d(scale_factor=2)
        self.act = nn.Tanh()

        # facelandmark
        self.fconv4_1 = InvertedResidual(64,128,2,False,2)

        # 第5层卷积层，Bottleneck结构
        # 重复6次
        self.fconv5_1 = InvertedResidual(128,128,1,False,4)
        self.fblock5_2 = InvertedResidual(128,128,1,True,4)
        self.fblock5_3 = InvertedResidual(128,128,1,True,4)
        self.fblock5_4 = InvertedResidual(128,128,1,True,4)
        self.fblock5_5 = InvertedResidual(128,128,1,True,4)
        self.fblock5_6 = InvertedResidual(128,128,1,True,4)

        # 第6层卷积层，Bottleneck结构
        # 重复1次
        self.fconv6_1 = InvertedResidual(128,16,1,False,2)

        # 第7层卷积层，重复1次
        self.fconv7 = conv_bn(16,32,3,2)
        # 第8层卷积层，重复1次
        self.fconv8 = nn.Conv2d(32,64,7,1,0)
        self.fbn8 = nn.BatchNorm2d(64)


        self.favg_pool1 = nn.AvgPool2d(14)
        self.favg_pool2 = nn.AvgPool2d(7)
        self.favg_pool3 = nn.AvgPool2d(5)

        self.ffc = nn.Linear(448,196)

        self.frelu = nn.ReLU(inplace=True)
        
        
        
        
        self.seresblock = nn.Sequential(*[ResBlock(channel * 4, channel * 4) for i in range(num_blocks)])

        self.seconv5 = nn.Conv2d(channel * 4, channel * 2, [3, 3], padding=1)  # 64,64
        self.sebn5 = nn.BatchNorm2d(channel * 2)
        self.seconv6 = nn.Conv2d(channel * 2, channel * 2, [3, 3], padding=1)  # 64,64
        self.sebn6 = nn.BatchNorm2d(channel * 2)
        self.seconv7 = nn.Conv2d(channel * 2, channel, [3, 3], padding=1)  # 64,64
        self.sebn7 = nn.BatchNorm2d(channel)
        self.seconv8 = nn.Conv2d(channel, channel, [3, 3], padding=1)  # 64,64
        self.sebn8 = nn.BatchNorm2d(channel)
        self.seconv9 = nn.Conv2d(channel, n_class, [7, 7], padding=3)  # 64,64
        self.relu = nn.ReLU(inplace=True)
        
    
    def forward(self, mem):
        landmark = None
        gan = None
        seg = None
        if 'facelandmark' in mem.keys():
            inputs_facelandmark = mem['facelandmark']['data']
            landmark = self.forward_facelandmark(inputs_facelandmark)
        if 'gan' in mem.keys():
            inputs_gan = mem['gan']['data']
            gan = self.forward_gan(inputs_gan)
        if 'seg' in mem.keys():
            inputs_seg = mem['seg']['data']   
            seg = self.forward_seg(inputs_seg)
            
        
        return gan,landmark,seg
    def forward_facelandmark(self,inputs):
        x0 = self.conv(inputs)
        x0 = self.bn(x0)
        x0 = self.leak_relu(x0)  # 256, 256, 32

        x1 = self.conv1(x0)
        x1 = self.leak_relu(x1)
        x1 = self.bn1(x1)
        x1 = self.conv2(x1)
        x1 = self.leak_relu(x1)  # 128, 128, 64'
        x1 = self.bn2(x1)

        x2 = self.conv3(x1)
        x2 = self.leak_relu(x2)
        x2 = self.bn3(x2)
        x2 = self.conv4(x2)
        x2 = self.leak_relu(x2)  # 64, 64, 128
        x2 = self.bn4(x2)
        
        x41=self.fconv4_1(x2)
        x51=self.fconv5_1(x41)
        x52=self.fblock5_2(x51)
        x53=self.fblock5_3(x52)
        x54=self.fblock5_4(x53)
        x55=self.fblock5_5(x54)
        x56=self.fblock5_6(x55)
        
        x61=self.fconv6_1(x56)
        
        output1= self.favg_pool1(x61)
        output1 = output1.view(output1.size(0),-1)
        
        x7=self.fconv7(x61)
        output2 = self.favg_pool2(x7)
        output2 = output2.view(output2.size(0),-1)
        output3 = self.favg_pool3(self.frelu(self.fconv8(x7)))
        output3 = output3.view(output3.size(0),-1)

        multi_scale = torch.cat([output1,output2,output3],1)
        labdmarks = self.ffc(multi_scale)
        
        return labdmarks
    
    def forward_gan(self,inputs):
        
        x0 = self.conv(inputs)
        x0 = self.bn(x0)
        x0 = self.leak_relu(x0)  # 256, 256, 32

        x1 = self.conv1(x0)
        x1 = self.leak_relu(x1)
        x1 = self.bn1(x1)
        x1 = self.conv2(x1)
        x1 = self.leak_relu(x1)  # 128, 128, 64'
        x1 = self.bn2(x1)

        x2 = self.conv3(x1)
        x2 = self.leak_relu(x2)
        x2 = self.bn3(x2)
        x2 = self.conv4(x2)
        x2 = self.leak_relu(x2)  # 64, 64, 128
        x2 = self.bn4(x2)



        x2 = self.resblock(x2)  # 64, 64, 128
        x2 = self.conv5(x2)
        x2 = self.leak_relu(x2)  # 64, 64, 64
        x2 = self.bn5(x2)

        x3 = self.upsample(x2)
        x3 = self.conv6(x3 + x1)
        x3 = self.leak_relu(x3)
        x3 = self.bn6(x3)
        x3 = self.conv7(x3)
        x3 = self.leak_relu(x3)  # 128, 128, 32
        x3 = self.bn7(x3)

        x4 = self.upsample(x3)
        x4 = self.conv8(x4 + x0)
        x4 = self.leak_relu(x4)
        x4 = self.bn8(x4)
        x4 = self.conv9(x4)  # 256, 256, 32
        x4 = self.bn9(x4)

        result = self.act(x4)

        return  torch.clamp(result,-0.999999,0.999999)
    
    def forward_seg(self,inputs):
        x0 = self.conv(inputs)
        x0 = self.bn(x0)
        x0 = self.leak_relu(x0)  # 256, 256, 32

        x1 = self.conv1(x0)
        x1 = self.leak_relu(x1)
        x1 = self.bn1(x1)
        x1 = self.conv2(x1)
        x1 = self.leak_relu(x1)  # 128, 128, 64'
        x1 = self.bn2(x1)

        x2 = self.conv3(x1)
        x2 = self.leak_relu(x2)
        x2 = self.bn3(x2)
        x2 = self.conv4(x2)
        x2 = self.leak_relu(x2)  # 64, 64, 128
        x2 = self.bn4(x2)
        
        x2 = self.seresblock(x2)  # 64, 64, 128
        x2 = self.seconv5(x2)
        x2 = self.relu(x2)  # 64, 64, 64
        x2 = self.sebn5(x2)

        x3 = self.upsample(x2)
        x3 = self.seconv6(x3 + x1)
        x3 = self.relu(x3)
        x3 = self.sebn6(x3)
        x3 = self.seconv7(x3)
        x3 = self.relu(x3)  # 128, 128, 32
        x3 = self.sebn7(x3)

        x4 = self.upsample(x3)
        x4 = self.seconv8(x4 + x0)
        x4 = self.relu(x4)
        x4 = self.sebn8(x4)
        x4 = self.seconv9(x4)  # 256, 256, 32
        
        return x4